/*
 * Copyright The OpenTelemetry Authors
 * SPDX-License-Identifier: Apache-2.0
 */

package com.xiashitech.agent.instrumentation.redis.redisson;

import com.xiashitech.agent.config.StartConfig;
import com.xiashitech.agent.instrumentation.redis.RedisData;
import com.xiashitech.agent.instrumentation.redis.constants.CommandNames;
import com.xiashitech.agent.utils.SessionUtil;
import io.netty.buffer.ByteBuf;
import io.opentelemetry.javaagent.bootstrap.internal.CommonConfig;
import io.opentelemetry.javaagent.shaded.instrumentation.api.db.RedisCommandSanitizer;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.baggage.Baggage;
import lombok.Getter;
import lombok.Setter;
import org.redisson.client.protocol.CommandData;
import org.redisson.client.protocol.CommandsData;

import javax.annotation.Nullable;
import java.lang.invoke.MethodHandle;
import java.lang.invoke.MethodHandles;
import java.lang.invoke.MethodType;
import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import static java.util.Collections.emptyList;
import static java.util.Collections.singletonList;

@Getter
@Setter
public class XSRedissonRequest {

    private Object command;

    private Baggage baggage;

    private static final RedisCommandSanitizer sanitizer =
            RedisCommandSanitizer.create(CommonConfig.get().isStatementSanitizationEnabled());

    /**
     * 设置请求参数
     *
     * @param redisData
     */
    public void setRedisDataParams(RedisData redisData) {
        String operation = getOperation();
        if (operation == null || "".equals(operation)) {
            return;
        }
        List<Object> params = getParamForOperation(operation);
        if (params == null || params.size() == 0) {
            return;
        }
        if (CommandNames.GET.getName().equals(operation)) {
            Object param0 = params.get(0);
            if (param0 instanceof byte[]) {
            }
            redisData.setKey(SessionUtil.convertObjectToByteArray(param0));
            if (param0 instanceof byte[]) {
            }
        } else if (CommandNames.SETEX.getName().equals(operation) || CommandNames.PSETEX.getName().equals(operation)) {
            if (params.size() < 3) {
                return;
            }
            Object param0 = params.get(0);
            Object param1 = params.get(1);
            Object param2 = params.get(2);
            redisData.setKey(SessionUtil.convertObjectToByteArray(param0));
            redisData.setTime(Long.valueOf(String.valueOf(param1)));
            redisData.setValue(SessionUtil.convertObjectToByteArray(param2));
        }
    }

    public List<Object> getParamForOperation(String operation) {
        if (operation == null) {
            return new ArrayList();
        }
        List<Object> params = new ArrayList();
        Object command = getCommand();
        if (command instanceof CommandData) {
            params.addAll(Arrays.asList(((CommandData<?, ?>) command).getParams()));
        } else if (command instanceof CommandsData) {
            CommandsData commandsData = (CommandsData) command;
            Optional<CommandData<?, ?>> commandDataOptional = commandsData.getCommands().stream().filter(commandData -> commandData.getCommand().getName().equals(operation)).findFirst();
            if (commandDataOptional.isPresent()) {
                params.addAll(Arrays.asList(commandDataOptional.get().getParams()));
            }
        }
        return params;
    }

    public boolean isSupportedCommand() {
        Object command = getCommand();
        if (command == null) {
            return false;
        }
        if (command instanceof CommandData) {
            CommandData commandData = (CommandData) command;
            if (!CommandNames.isSupportedCommand(commandData.getCommand().getName())) {
                return false;
            }
        } else if (command instanceof CommandsData) {
            CommandsData commandsData = (CommandsData) command;
            if (!commandsData.getCommands().stream().anyMatch(commandData -> CommandNames.isSupportedCommand(commandData.getCommand().getName()))) {
                return false;
            }
        } else {
            return false;
        }
        return true;
    }

    public String getOperation() {
        Optional<String> commandNameOptional = getOperations().stream().filter(commandName -> CommandNames.getALLCommandName().contains(commandName)).findFirst();
        if (commandNameOptional.isPresent()) {
            return commandNameOptional.get();
        } else {
            return null;
        }
    }

    @Nullable
    public List<String> getOperations() {
        List<String> commonNames = new ArrayList();
        Object command = getCommand();
        if (command instanceof CommandData) {
            commonNames.add(((CommandData<?, ?>) command).getCommand().getName());
        } else if (command instanceof CommandsData) {
            CommandsData commandsData = (CommandsData) command;
            commonNames.addAll(commandsData.getCommands().stream().map(commandData -> commandData.getCommand().getName()).collect(Collectors.toList()));
        }
        return commonNames;
    }

    public List<Object> getALLParams() {
        List<Object> params = new ArrayList();
        Object command = getCommand();
        if (command instanceof CommandData) {
            params.addAll(Arrays.asList(((CommandData<?, ?>) command).getParams()));
        } else if (command instanceof CommandsData) {
            CommandsData commandsData = (CommandsData) command;
            params.addAll(commandsData.getCommands().stream().map(commandData -> commandData.getParams()).flatMap(paramArray -> Stream.of(paramArray)).collect(Collectors.toList()));
        }
        return params;
    }

    @Nullable
    public String getStatement() {
        List<String> sanitizedStatements = sanitizeStatement();
        switch (sanitizedStatements.size()) {
            case 0:
                return null;
            // optimize for the most common case
            case 1:
                return sanitizedStatements.get(0);
            default:
                return String.join(";", sanitizedStatements);
        }
    }

    private List<String> sanitizeStatement() {
        Object command = getCommand();
        // get command
        if (command instanceof CommandsData) {
            List<CommandData<?, ?>> commands = ((CommandsData) command).getCommands();
            return commands.stream()
                    .map(XSRedissonRequest::normalizeSingleCommand)
                    .collect(Collectors.toList());
        } else if (command instanceof CommandData) {
            return singletonList(normalizeSingleCommand((CommandData<?, ?>) command));
        }
        return emptyList();
    }

    public static String normalizeSingleCommand(CommandData<?, ?> command) {
        Object[] commandParams = command.getParams();
        List<Object> args = new ArrayList<>(commandParams.length + 1);
        if (command.getCommand().getSubName() != null) {
            args.add(command.getCommand().getSubName());
        }
        for (Object param : commandParams) {
            if (param instanceof ByteBuf) {
                try {
                    // slice() does not copy the actual byte buffer, it only returns a readable/writable
                    // "view" of the original buffer (i.e. read and write marks are not shared)
                    ByteBuf buf = ((ByteBuf) param).slice();
                    // state can be null here: no Decoders used by Codecs use it
                    args.add(command.getCodec().getValueDecoder().decode(buf, null));
                } catch (Exception ignored) {
                    args.add("?");
                }
            } else {
                args.add(param);
            }
        }
        return sanitizer.sanitize(command.getCommand().getName(), args);
    }

    @Nullable
    public XSPromiseWrapper<?> getPromiseWrapper() {
        CompletionStage<?> promise = getPromise();
        if (promise instanceof XSCompletableFuture<?>) {
            return (XSPromiseWrapper<?>) promise;
        }
        return null;
    }

    @Nullable
    private CompletionStage<?> getPromise() {
        Object command = getCommand();
        if (command instanceof CommandData && COMMAND_DATA_GET_PROMISE != null) {
            try {
                return (CompletionStage<?>) COMMAND_DATA_GET_PROMISE.invoke(command);
            } catch (Throwable ignored) {
                StartConfig.throwableStackPrinter(ignored);
                return null;
            }
        } else if (command instanceof CommandsData && COMMANDS_DATA_GET_PROMISE != null) {
            try {
                return (CompletionStage<?>) COMMANDS_DATA_GET_PROMISE.invoke(command);
            } catch (Throwable ignored) {
                StartConfig.throwableStackPrinter(ignored);
                return null;
            }
        }
        return null;
    }

    private static final MethodHandle COMMAND_DATA_GET_PROMISE =
            findGetPromiseMethod(CommandData.class);
    private static final MethodHandle COMMANDS_DATA_GET_PROMISE =
            findGetPromiseMethod(CommandsData.class);

    private static MethodHandle findGetPromiseMethod(Class<?> commandClass) {
        MethodHandles.Lookup lookup = MethodHandles.publicLookup();
        try {
            Class<?> promiseClass =
                    Class.forName(
                            "org.redisson.misc.RPromise", false, XSRedissonRequest.class.getClassLoader());
            // try versions older than 3.16.8
            return lookup.findVirtual(commandClass, "getPromise", MethodType.methodType(promiseClass));
        } catch (NoSuchMethodException | ClassNotFoundException e) {
            // in 3.16.8 CommandsData#getPromise() and CommandData#getPromise() return type was changed in
            // a backwards-incompatible way from RPromise to CompletableFuture
            try {
                return lookup.findVirtual(
                        commandClass, "getPromise", MethodType.methodType(CompletableFuture.class));
            } catch (NoSuchMethodException | IllegalAccessException ignored) {
                return null;
            }
        } catch (IllegalAccessException ignored) {
            return null;
        }
    }
}
